# -*- coding: utf-8 -*-
# @Time    : 2024/04/24 08:41
# @Author  : Mr.su
# @FileName: download_image.py
# @FileDesc: 图片下载功能
from queue import Queue
from CollectSpiders.toots.constants import WEBTYPE_DICT
from CollectSpiders.settings import MONGODB_CONFIG, BEANSTALK_CONFIG
import requests, pymongo, time, json, threading, datetime, os, beanstalkc
from CollectSpiders.toots.methods import make_md5, make_img_name, random_str


class DownLoadImage:
    headers = {
        'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36',
    }

    def __init__(self):
        self.products = None
        self.images = None
        self.img_queue = Queue()
        self.path = "./data/work/new_rule/{}/{}-{}/{}/{}/{}/"

    def download(self, img_dic, path):
        if not os.path.exists(path):
            os.makedirs(path)
        save_path = '{}{}'.format(path, img_dic['name'])
        try:
            res = requests.get(img_dic['url'], headers=self.headers, timeout=20)
            with open(save_path, 'wb') as f:
                f.write(res.content)
            print('{} [info]: 下载图片 {}'.format(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), save_path))
            return save_path, True
        except Exception as e:
            print('{} [error]: {} 图片{}下载失败'.format(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), e, img_dic['_id']))
            return save_path, False

    def process_product(self, product):
        current = datetime.datetime.now().strftime('%Y-%m-%d')
        num = 0
        img_path_lis = []
        add_status = True
        for img_url in product['image']:
            estr = random_str(8)
            path = self.path.format(WEBTYPE_DICT[product['webtype']], product['domain'], product['pid'], current, product['label'], product['brand'])
            img_dic = {
                "_id": make_md5(img_url), "_class": "spider.resource.model.DownloadedLink",
                "product_id": product['_id'], "pid": product['pid'], "site_id": product['sid'],
                "url": img_url, "add_time": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                "path": path, "copy_path": "", "name": ""
            }
            if product['image'].index(img_url) == 0:
                img_dic['first'] = 'true'
                img_dic['name'] = make_img_name(product['name'], estr, product['brand'], True)
            else:
                img_dic['name'] = make_img_name(product['name'], estr, product['brand'], False, num)
            num += 1
            result_lis = []
            for i in range(2, 7):
                save_path, result = self.download(img_dic, img_dic['path'])
                if result:
                    self.images.update_one({'_id': img_dic['_id']}, {'$setOnInsert': img_dic}, upsert=True)
                    img_path_lis.append(img_dic['path'] + img_dic['name'])
                    break
                result_lis.append(result)
                self.products.update_one({'_id': product['_id']}, {"$set": {"status": "{}".format(i)}})
            if len(result_lis) == 5:
                add_status = False
                break
        if add_status:
            self.products.update_one({'_id': product['_id']}, {"$set": {"status": "1"}})
            print('{} [info]: 当前id({})商品全部图片下载完成...'.format(datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'), product['_id']))
        # TODO 商品图片全部下载完成后上传功能
        # if product['upload'] == 1:
        #     for img_path in img_path_lis:
        #         data_dic = json.dumps({
        #             "document_id": product_dic['_id'], "file_path": img_path,
        #             "article_id": 0  # 如果没有ID就设置为0
        #         })
        #         msg = beanstalk_conn.put(data_dic)

    def download_start(self):
        while True:
            if self.img_queue.empty():
                time.sleep(60)
                continue
            product = json.loads(self.img_queue.get())
            t = threading.Thread(target=self.process_product, args=(product,))
            t.start()
            self.img_queue.task_done()
            time.sleep(2)

    def run(self):
        t_download = threading.Thread(target=self.download_start)
        t_download.start()

        while True:
            # 链接mongo
            conn = pymongo.MongoClient(MONGODB_CONFIG)
            self.products = conn["wow_spider_db"]['tb_product']
            self.images = conn["wow_spider_url_db"]['download_links']
            query = {'status': '0'}
            results = self.products.find(query)
            for result in results:
                self.img_queue.put(json.dumps(result, ensure_ascii=False))
            time.sleep(1800)


if __name__ == '__main__':
    # BEANSTALK_QUEUE = "wow_upload_pic_task"
    # beanstalk_conn = beanstalkc.Connection(**BEANSTALK_CONFIG)
    # beanstalk_conn.use(BEANSTALK_QUEUE)
    # beanstalk_conn.using()
    d = DownLoadImage()
    d.run()
